
from tensorflow.keras.models import load_model

from bert4keras.tokenizers import Tokenizer,load_vocab

#因为需要加载custom_object
from bert4keras.layers import *


from utils.bert_info import BertInfo
from utils.generator import BertConditionGenerator
max_len = 128

model_saved_path = './model_saved/best_model'

bert_info_obj = BertInfo()

token_dict,keep_tokens = load_vocab(
    dict_path=bert_info_obj.dict_path,
    simplified=True,
    startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]'],
)
tokenizer = Tokenizer(token_dict,do_lower_case=True)

generate_model = load_model(model_saved_path)

sentiment_generator = BertConditionGenerator(
    start_id=tokenizer._token_start_id,
    end_id = tokenizer._token_end_id,
    maxlen=max_len,
    tokenizer=tokenizer,
    generator_model=generate_model
)


def just_show():
    print("正面采样:")
    print(sentiment_generator.generate(1,5,5),"\n")
    print("负面采样")
    print(sentiment_generator.generate(0,5,5),"\n")


if __name__ == '__main__':
    just_show()